import os
import os.path as osp
import numpy as np
import cv2
import re

from tqdm import tqdm # For extracting file IDs

# Assuming BaseMultiViewDataset and utility functions are in these locations
# Adjust the import paths based on your project structure
from src.dust3r.datasets.base.base_multiview_dataset import BaseMultiViewDataset, cropping
from src.dust3r.utils.image import imread_cv2
# from dust3r.utils.some_module import efficient_random_intervals_revised # Placeholder for sampling logic

# If efficient_random_intervals_revised is not available or you prefer Co3D's sampling:
# You might need to implement or adapt a sampling function.
# For this example, I'll include a placeholder for where it would be called.
# BaseMultiViewDataset likely has a method like 'get_seq_from_start_id'
# which could be an alternative if 'efficient_random_intervals_revised' is specific
# to ScanNetpp_Multi's internal structure and not general.

class WildRGBD_Multi(BaseMultiViewDataset):
    def __init__(self, *args, ROOT, samples_per_scene=1, max_interval=5,
                 mask_bg="rand", # True, False, or "rand" like in Co3d_Multi
                 check=False,
                 **kwargs):
        self.ROOT = ROOT
        self.video = True  # Assuming sequential data
        self.is_metric = True
        self.max_interval = max_interval
        self.samples_per_scene = samples_per_scene
        self.dataset_label = "WildRGBD"
        self.check = check
        
        assert mask_bg in (True, False, "rand")
        self.mask_bg = mask_bg

        super().__init__(*args, **kwargs) # Initializes self.num_views, self.allow_repeat etc.

        self._load_data()

    def _load_data(self):
        self.all_subscenes = []  # List of (scene_name, subscene_name) tuples
        self.subscene_data = {}  # Dict: (scene, subscene) -> sorted list of fileids

        for scene_name in sorted(os.listdir(self.ROOT)):
            scene_path = osp.join(self.ROOT, scene_name)
            if not osp.isdir(scene_path):
                continue

            for subscene_name in sorted(os.listdir(scene_path)):
                subscene_path = osp.join(scene_path, subscene_name)
                if not osp.isdir(subscene_path):
                    continue

                current_fileids = []
                for f_name in os.listdir(subscene_path):
                    if f_name.endswith("_rgb.png"):
                        file_id = f_name.removesuffix("_rgb.png")
                        # Check if corresponding depth and cam data exist
                        depth_path = osp.join(subscene_path, f"{file_id}_depth.png")
                        cam_path = osp.join(subscene_path, f"{file_id}_cam.npz")
                        if osp.exists(depth_path) and osp.exists(cam_path):
                            current_fileids.append(file_id)

                if not current_fileids:
                    if not os.listdir(subscene_path):
                        os.rmdir(subscene_path)
                        print(f"Successfully remove the empty directory: {subscene_path}")
                    print(f"Skipping {subscene_path}: No valid file_ids with rgb, depth, and cam.npz found.")
                    continue

                # Sort the file IDs to ensure a consistent order
                sorted_unique_fileids = sorted(current_fileids)
                
                if not self.check:
                    # Ensure enough views for a sequence
                    cut_off = (
                        self.num_views if not self.allow_repeat else max(self.num_views // 3, 3)
                    )
                    if len(sorted_unique_fileids) < cut_off:
                        print(f"Skipping {subscene_path}: Only {len(sorted_unique_fileids)} valid frames, needs {cut_off}.")
                        continue
                
                subscene_tuple = (scene_name, subscene_name)
                self.all_subscenes.append(subscene_tuple)
                self.subscene_data[subscene_tuple] = sorted_unique_fileids
        
        if not self.all_subscenes:
            raise RuntimeError(f"No valid subscenes found in {self.ROOT}. Please check directory structure and file names (e.g., 000100_rgb.png, 000100_depth.png, 000100_cam.npz).")

    def __len__(self):
        return len(self.all_subscenes) * self.samples_per_scene

    def check_depth_valid(self, actual_mask_bg=True, resolution=(224, 224)):
        outdir = osp.join(osp.dirname(self.ROOT), 'discarded')
        # os.makedirs(outdir, exist_ok=True)
        import shutil
        import glob
        for scene_name, subscene_name in tqdm(self.all_subscenes):
            available_fileids = self.subscene_data[(scene_name, subscene_name)]
            subscene_dir = osp.join(self.ROOT, scene_name, subscene_name)
            dest_dir = osp.join(outdir, scene_name)
            if len(available_fileids) < 10:
                print(f"views of {subscene_dir} < 10 to {dest_dir}")
                shutil.move(subscene_dir, dest_dir)
                if osp.isdir(subscene_dir) and not os.listdir(subscene_dir):
                    os.rmdir(subscene_dir)
                    print(f"Successfully remove the empty directory: {subscene_dir}")
                continue

            for file_id in available_fileids:
                rgb_path = osp.join(subscene_dir, f"{file_id}_rgb.png")
                depth_path = osp.join(subscene_dir, f"{file_id}_depth.png")
                mask_path = osp.join(subscene_dir, f"{file_id}_mask.png") # For object mask
                cam_path = osp.join(subscene_dir, f"{file_id}_cam.npz")   # Camera parameters
                depthmap = imread_cv2(depth_path, cv2.IMREAD_UNCHANGED)
                depthmap[~np.isfinite(depthmap)] = 0  # invalid

                if actual_mask_bg:
                    object_mask = imread_cv2(mask_path, cv2.IMREAD_UNCHANGED)
                    object_mask = (object_mask.astype(np.float32) / 255.0) > 0.1 # Example threshold
                    
                    depthmap *= object_mask

                cam_data = np.load(cam_path)
                intrinsics = cam_data["intrinsics"].astype(np.float32) # Ensure key matches your npz file
                rgb_image = imread_cv2(rgb_path, cv2.IMREAD_COLOR)
                rgb_image, depthmap, intrinsics = self._crop_resize_if_necessary(
                    rgb_image, depthmap, intrinsics, resolution, info=f"{file_id}"
                )

                num_valid_depth = (depthmap > 0.0).sum() # Check for valid depth after processing
                if num_valid_depth < 50:
                     print(f"Warning: {num_valid_depth} valid depth for {file_id} in {subscene_dir} ")
                     os.makedirs(dest_dir, exist_ok=True)
                     for file_to_move in glob.glob(osp.join(subscene_dir, f"{file_id}*")):
                        try:
                            shutil.move(file_to_move, dest_dir)
                        except shutil.Error as e:
                            print(f"Error moving file {file_to_move}: {e}")


    def _get_views(self, idx, resolution, rng, num_views):
        subscene_tuple_idx = idx // self.samples_per_scene
        scene_name, subscene_name = self.all_subscenes[subscene_tuple_idx]
        
        available_fileids = self.subscene_data[(scene_name, subscene_name)]
        
        pos, ordered_video = self.efficient_random_intervals_revised(
            0,
            len(available_fileids),
            num_views,
            rng,
            min_interval=1,
            max_interval=self.max_interval,
        )
        chosen_fileids = np.array(available_fileids)[pos]

        views = []
        subscene_dir = osp.join(self.ROOT, scene_name, subscene_name)

        # Determine if background should be masked for this sample
        actual_mask_bg = (self.mask_bg == True) or \
                         (self.mask_bg == "rand" and rng.choice([True, False], p=[0.1, 0.9]))


        for i, file_id in enumerate(chosen_fileids):
            rgb_path = osp.join(subscene_dir, f"{file_id}_rgb.png")
            depth_path = osp.join(subscene_dir, f"{file_id}_depth.png")
            mask_path = osp.join(subscene_dir, f"{file_id}_mask.png") # For object mask
            cam_path = osp.join(subscene_dir, f"{file_id}_cam.npz")   # Camera parameters

            if not all(osp.exists(p) for p in [rgb_path, depth_path, cam_path]):
                # This should ideally not happen if _load_data is thorough
                # Or handle by retrying / skipping, similar to Co3d_Multi's invalidation
                print(f"Warning: Missing file for {file_id} in {subscene_dir}. Skipping view.")
                # Needs more robust handling for production: either ensure files exist or implement retry/invalidation
                continue 

            rgb_image = imread_cv2(rgb_path, cv2.IMREAD_COLOR)
            # Original WildRGBD depth processing
            depthmap = imread_cv2(depth_path, cv2.IMREAD_UNCHANGED)
            depthmap = depthmap.astype(np.float32) / 1000.0 
            depthmap[~np.isfinite(depthmap)] = 0  # invalid

            cam_data = np.load(cam_path)
            intrinsics = cam_data["intrinsics"].astype(np.float32) # Ensure key matches your npz file
            camera_pose = cam_data["pose"].astype(np.float32)   # Ensure key matches your npz file

            if actual_mask_bg:
                object_mask = imread_cv2(mask_path, cv2.IMREAD_UNCHANGED)
                object_mask = (object_mask.astype(np.float32) / 255.0) > 0.1 # Example threshold
                
                depthmap *= object_mask
            
            # Placeholder: If your depth has specific invalid values like ScanNet++ (65535)
            # depthmap[depthmap == SPECIFIC_INVALID_VALUE] = 0 

            rgb_image, depthmap, intrinsics = self._crop_resize_if_necessary(
                rgb_image, depthmap, intrinsics, resolution, rng=rng, info=f"{file_id}"
            )

            intrinsics, camera_pose = cropping.get_center_camera(intrinsics, camera_pose, depthmap=depthmap)
            # num_valid_depth = (depthmap > 0.0).sum() # Check for valid depth after processing
            # if num_valid_depth == 0:
            #      # Handle invalid image (e.g. all depth is zero after masking/cropping)
            #      # This part would need the invalidation logic from Co3d_Multi if it's a common issue
            #      print(f"Warning: Zero valid depth for {file_id} in {subscene_dir} after processing. Skipping view.")
            #      # A robust way would be to retry sampling for this 'idx' or mark image as invalid.
            #      # For simplicity, we continue, but this might lead to fewer than num_views.
            #      continue


            img_mask, ray_mask = self.get_img_and_ray_masks(
                self.is_metric, len(views), rng # len(views) is current view index in sequence
            )

            views.append(
                dict(
                    img=rgb_image,
                    depthmap=depthmap.astype(np.float32),
                    camera_pose=camera_pose.astype(np.float32),
                    camera_intrinsics=intrinsics.astype(np.float32),
                    dataset='wildrgbd',
                    label=f"{scene_name}/{subscene_name}/{file_id}",
                    instance=f"{scene_name}_{subscene_name}_{file_id}", # More descriptive instance
                    is_metric=self.is_metric,
                    is_video=ordered_video, # From get_seq_from_start_id
                    quantile=np.array(0.85, dtype=np.float32), # Like original WildRGBD
                    img_mask=img_mask,
                    ray_mask=ray_mask,
                    # Fields from BaseMultiViewDataset/Co3d_Multi
                    camera_only=False,
                    depth_only=False,
                    single_view=False,
                    reset=False, 
                )
            )
        assert len(views) == num_views
        return views